from locale import normalize
import numpy as np
from scipy.spatial import cKDTree
from matplotlib import pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator
import pcl
import pandas

import csv
import cv2
import os, sys

def sort_by_angle(data, angle_resolution):
    incidence = np.uint32((data[:, 3] + 180) / angle_resolution)
    # incidence = data[:, 3]
    by_incidence = { i: [] for i in range(incidence.min(), incidence.max() + 1)}
    for i in range(data.shape[0]):
        if data[i, 4] == 0: continue
        by_incidence[incidence[i]].append(i)
        data[i, 0] = data[i, 0] - data[by_incidence[incidence[i]][0], 0]

    return by_incidence
    

def read_scan(path):
    with open(path, 'r') as f:
        f_csv = csv.reader(f)
        scan_data = [[float(item) for item in row] for row in f_csv]
        return np.float64(scan_data)

def segment_all_lines(full_cloud):    
    cloud = full_cloud
    lines=[]
    while(cloud.size>100):
        segmenter = cloud.make_segmenter()
        segmenter.set_optimize_coefficients(True)
        segmenter.set_model_type(pcl.SACMODEL_LINE)
        segmenter.set_method_type(pcl.SAC_RANSAC)
        segmenter.set_distance_threshold(100)
        segmenter.set_MaxIterations(50)
        indices, coeff = segmenter.segment()
        if len(indices) == 0:
            break

        coeff = np.array(coeff)
        point_on_line = coeff[:3]
        line_dir = coeff[3:]
        perp_dir = np.cross(line_dir, [0, 0, 1])        
        perp_dir = perp_dir / magtitude(perp_dir)
        if np.dot(point_on_line, perp_dir) > 0:
            perp_dir = -perp_dir
        lines.append([len(indices), point_on_line, line_dir, perp_dir])
        cloud = cloud.extract(indices, negative=True)

    lines.sort(key=lambda line: line[0], reverse=True)
    return lines

def magtitude(data):
    return np.sqrt(np.sum(np.square(data), axis=-1))

if __name__ == "__main__":
    # # 温度，水平转角，竖直转角，距离，反射率，cos
    # for f in os.listdir('./data/fixed'):
    #     fname = os.path.splitext(f)
    #     if fname[1] == '.csv':
    #         data = read_scan('./data/fixed/%s' % f)            
    #         np.save('./data/fixed/%s.npy' % fname[0], data)
    # data = read_scan('./data/fixed.csv')
    # np.save('./data/fixed.npy', data)
    # data = np.load('./data/fixed/cell_leftMiddle_fix_20220121_1722.npy')
    # data = np.load('./data/fixed/cell_rightClose_fix_20220121_1713.npy')
    # data = np.load('./data/fixed/cell_rightClose_fix_20220121_1707.npy')
    file =  './data/fixed/0124/cell_fix_test2_20220124_1912.csv'
    data = pandas.read_csv(file, header=None,
        names=['timestamp', 'temperature', 'hangle', 'vangle', 'dist', 'intensity', 'x', 'y', 'z', 'iangle', 'plane_offset'],
        dtype={'timestamp': np.float64, 'temperature': np.int32, 'hangle': np.float32, 'vangle': np.float32, 'dist': np.float32
        , 'intensity': np.float32, 'x': np.float32, 'y': np.float32, 'z': np.float32
         })
    data = np.array(data)
    data[data[:, 4] < 300] = 0
    data[data[:, 4] > 10000] = 0
    # data = data_all[data_all[:, 3] < np.pi, :]    
    res = 0.05
    by_incidence = sort_by_angle(data, res)
    avg_rho = np.zeros((int(360 / res + 1), 4), dtype=np.float64)
    for k in by_incidence.keys():
        idx = by_incidence[k]
        if len(idx) == 0: continue
        rhos_by_angle = data[idx, 4]
        avg_rho[k, 0] = rhos_by_angle.mean()
        avg_rho[k, 1] = rhos_by_angle.std()
        avg_rho[k, 2] = avg_rho[k, 1] / avg_rho[k, 0]
        avg_rho[k, 3] = avg_rho[k, 1] * avg_rho[k, 0]

    # f = plt.figure('angle', figsize=(20, 2))
    # i = 0
    # for k in by_incidence.keys():
    #     idx = by_incidence[k]
    #     if len(idx) == 0: continue
    #     plt.plot(data[idx, 0][:50], data[idx, 4][:50] - avg_rho[k, 0])
    #     i = i + 1
    #     if i > 2:
    #         break
    # plt.show()

    # f = plt.figure('avg std', figsize=(20, 5))
    angles = np.arange(avg_rho.shape[0]) * res - 180
    idx = np.argwhere((angles > -150) * (angles < 150))
    # plt.plot(angles[idx], avg_rho[idx, 1])
    # plt.show()

    # f = plt.figure('relative std', figsize=(20, 5))
    # plt.plot(angles[idx], avg_rho[idx, 2])
    # plt.show()

    # f = plt.figure('var', figsize=(20, 5))
    # plt.plot(angles[idx], avg_rho[idx, 3])
    
    # f = plt.figure('polar', figsize=(10, 10))
    rads = np.deg2rad(angles)
    avg_xyz = np.zeros((rads.shape[0], 3), dtype=np.float32)
    avg_xyz[:, 0] = (np.sin(rads) * avg_rho[:, 0]).reshape(-1)
    avg_xyz[:, 1] = (np.cos(rads) * avg_rho[:, 0]).reshape(-1)
    # plt.plot(avg_xyz[:, 0], avg_xyz[:, 1], '.', label='avg polar')
    angle = [0, -30, -120, -150, 60, 120, 180]
    angle_idx = np.argwhere(np.isin(angles, angle))
    # rho_at_angle = avg_rho[angle_idx, 0]  
    # for aidx in angle_idx:
    #     plt.plot([0, float(np.sin(np.deg2rad(angles[aidx])) * avg_rho[aidx, 0])],
    #         [0, float(np.cos(np.deg2rad(angles[aidx])) * avg_rho[aidx, 0])], c='black')
    # angle = -50
    # rho_at_angle = avg_rho[np.argwhere(angles == angle), 0]
    # plt.plot(np.sin(np.deg2rad(angle)) * rho_at_angle, np.cos(np.deg2rad(angle)) * rho_at_angle, 'o',  c='black')
    # plt.show()    

    f = plt.figure('raw polar', figsize=(10, 15))
    ax = f.add_subplot(3, 1, (1, 2))
    plt.title(os.path.split(file)[-1])
    raw_angles = np.deg2rad(data[:, 3])
    xyz = np.zeros((raw_angles.shape[0], 3), dtype=np.float32)
    xyz[:, 0] = np.sin(raw_angles) * data[:, 4]
    xyz[:, 1] = np.cos(raw_angles) * data[:, 4]
    # lines = segment_all_lines(pcl.PointCloud(xyz))
    lines = segment_all_lines(pcl.PointCloud(avg_xyz))
    colors=['red', 'blue', 'green', 'orange']
    
    normalized_xyz = xyz.copy()
    mag_xyz = magtitude(xyz)
    normalized_xyz[:, 0] = normalized_xyz[:, 0] / mag_xyz
    normalized_xyz[:, 1] = normalized_xyz[:, 1] / mag_xyz
    normalized_xyz[:, 2] = normalized_xyz[:, 2] / mag_xyz
    for i in range(4):
        point_on_line = lines[i][1]
        line_dir = lines[i][2]
        perp_dir = lines[i][3]
        dists = np.dot(xyz - point_on_line, perp_dir)
        line_idx = np.abs(dists) < 20       
        ax.plot(xyz[line_idx, 0], xyz[line_idx, 1], '.', label='raw polar', c=colors[i])

    ax.plot(avg_xyz[:, 0], avg_xyz[:, 1], '.', label='avg polar', c='black')
    for aidx in angle_idx:
         ax.plot([0, float(np.sin(np.deg2rad(angles[aidx])) * avg_rho[aidx, 0])],
            [0, float(np.cos(np.deg2rad(angles[aidx])) * avg_rho[aidx, 0])], c='black')
    # plt.show()

    # plt.figure('incidence', figsize=(20,5))    
    ax = f.add_subplot(3, 1, 3)
    # lamdas = []
    # for i in range(4):
    #     point_on_line = lines[i][1] 
    #     line_dir = lines[i][2]
    #     perp_dir = lines[i][3]
    #     dists = np.dot(xyz - point_on_line, perp_dir)
    #     line_idx = np.abs(dists) < 20
    #     lamda_angle = np.rad2deg(np.arccos(np.abs(np.dot(normalized_xyz[line_idx], perp_dir))))
    #     lamdas.append([line_idx, dists[line_idx], lamda_angle])
    #     ax.plot(lamdas[i][2], lamdas[i][1], '.', c=colors[i])
    #     # plt.plot(lamda_angle, avg_rho[line_idx, 1], '.', c=colors[i])
    # ax.plot([0,90], [0, 0], c='black')

    # fig = plt.figure('avg incidence', figsize=(20,5))
    # ax = fig.add_subplot(projection='3d')

    normalized_avg_xyz = avg_xyz.copy()
    mag_avg_xyz = magtitude(avg_xyz)
    avg_xyz_nonzero = mag_avg_xyz > 0
    normalized_avg_xyz[avg_xyz_nonzero, 0] = normalized_avg_xyz[avg_xyz_nonzero, 0] / mag_avg_xyz[avg_xyz_nonzero]
    normalized_avg_xyz[avg_xyz_nonzero, 1] = normalized_avg_xyz[avg_xyz_nonzero, 1] / mag_avg_xyz[avg_xyz_nonzero]
    normalized_avg_xyz[avg_xyz_nonzero, 2] = normalized_avg_xyz[avg_xyz_nonzero, 2] / mag_avg_xyz[avg_xyz_nonzero]

    lamdas = []
    for i in range(4):
        point_on_line = lines[i][1]
        line_dir = lines[i][2]
        perp_dir = lines[i][3]
        dists = np.dot(avg_xyz - point_on_line, perp_dir)
        line_idx = np.abs(dists) < 20
        lamda_angle = np.rad2deg(np.arccos(np.abs(np.dot(normalized_avg_xyz[line_idx], perp_dir))))
        lamdas.append([line_idx, dists[line_idx], lamda_angle])
        ax.plot(lamdas[i][2], lamdas[i][1], '.', c=colors[i])
        # plt.plot(lamda_angle, avg_rho[line_idx, 1], '.', c=colors[i])
    ax.plot([0,90], [0, 0], c='black')
    plt.savefig(file+'.png')

    plt.figure('relative avg incidence', figsize=(20, 10))
    for i in range(4):
        plt.plot(lamdas[i][2], lamdas[i][1] / avg_rho[lamdas[i][0], 0], '.', c=colors[i])

    plt.plot([0,90], [0, 0], c='black')
    plt.show()

    plt.waitforbuttonpress()
    # fig, ax = plt.subplots(subplot_kw={"projection": "3d"})

    # # Make data.
    # time_step = 0
    # X = np.arange(data[:, -1].min(), data[:, -1].max(), res)
    # Y = np.arange(data[:, 0].min(), data[:, 0].max(), time_step)
    # X, Y = np.meshgrid(X, Y)
    # Z = np.zeros(X.shape, dtype=np.float64)


    # # Plot the surface.
    # surf = ax.plot_wireframe(X, Y, Z)

    # # Customize the z axis.
    # ax.set_zlim(-1.01, 1.01)
    # ax.zaxis.set_major_locator(LinearLocator(10))
    # # A StrMethodFormatter is used automatically
    # ax.zaxis.set_major_formatter('{x:.02f}')

    # # Add a color bar which maps values to colors.
    # fig.colorbar(surf, shrink=0.5, aspect=5)

    # plt.show()

    
